Skip to main content

01 - 数据工程

这一篇做完,磁盘上会多出一个 train.bin 文件。下一篇写的训练代码,打开它就能直接开训。

需要的前置知识:会写 Python,知道什么是数组和字典。不需要懂深度学习,也不需要读过 Transformer 的论文。涉及的概念我都会在用到的地方先解释一遍。

零、开始之前:模型到底吃什么

这一节不写代码,只建立一个心智模型。如果跳过这一节,后面每一步你都会觉得是在凭空做选择题。

0.1 模型不认识文字

神经网络内部全是矩阵乘法和加法,它处理的只能是数字。所以拿到一句中文或英文,第一件事必须是把它变成数字。

变成数字的方式不是随便编号。我们要的是:把文本切成一小片一小片,每一片在一张对照表里查到一个编号。这张对照表叫词表(vocabulary),切出来的每一片叫一个 token,查到的编号叫 token id

原文:   unhappiness
切片: un | happi | ness
查表: 1523 | 8891 | 2077

0.2 token 不是字,也不是词

初学最容易卡在这。token 是介于字符和单词之间的东西,中文里没有对应的概念,所以一般不翻译。

为什么不按单词切?因为单词太多了。英文常用词几十万,加上时态变化、复合词、人名地名,词表根本装不下。而且遇到词表里没有的词就完蛋了,只能标记成「未知」,信息直接丢失。

为什么不按字符切?因为那样序列太长了。unhappiness 有 11 个字符,就要占 11 个位置。模型每次能看的长度是有限的(我们这个模型是 2048),按字符切能装的内容就少得可怜,而且模型还得自己从字母学起单词的概念,学得非常慢。

所以取中间:常用的词保持完整,生僻的词拆成有意义的片段。unhappiness 拆成 un + happi + ness,三个片段都是高频的,都在词表里。这样词表规模可控(我们用 32000),又不会有装不下的词。

怎么决定哪些片段进词表,就是第四节要讲的 BPE 算法干的事。

0.3 训练的时候,一次输入长什么样

模型的训练目标只有一个:看着前面的 token,猜下一个 token 是什么。 这叫 next token prediction,也叫自回归。就这么简单,没有别的。

具体到一次输入。假设有一句话被切成了 5 个 token,id 是 [10, 24, 37, 42, 55]

输入 x = [10, 24, 37, 42]
标签 y = [24, 37, 42, 55]

y 就是 x 整体右移一位。模型看到 10,要预测 24;看到 10, 24,要预测 37;以此类推。四个位置同时算,同时算 loss,这就是为什么 Transformer 训练效率高。

模型每个位置输出的是一个长度 32000 的概率分布(词表里每个 token 各自的概率),拿它和真实的下一个 token 算交叉熵,就是 loss。

这里可以顺手记一个非常有用的数:训练刚开始、模型什么都还没学到的时候,它对 32000 个候选是均匀猜的,此时 loss 应该等于 ln(32000) ≈ 10.37。03 篇第一次跑训练时,如果初始 loss 不在 10.4 附近,说明代码有 bug,不用等它跑完。

0.4 所以数据工程要干的事

现在可以说清楚这一篇的任务了:把一大堆网页文本,变成一长串 token id,存成一个文件。

就这一件事。输入是几千万篇文档,输出是一个几十亿长度的整数数组。

"The cat sat..."  ┐
"机器学习是..." ├──► [10, 24, 37, 42, 55, 2, 891, 7734, ...] ──► train.bin
"Once upon a..." ┘

训练的时候,从这个长数组里随便挑一个位置,往后取 2048 个数当 x,右移一位取 2048 个数当 y,就是一条训练样本。

0.5 那为什么要五个步骤

既然任务这么简单,直接 tokenize 一下不就完了?中间那些清洗、去重是干嘛的。

因为原始网页数据很脏,而数据的问题没法在后面的环节补救,只能重训。模型结构写错了,loss 不降,两小时就能发现;数据里混了 5% 的重复网页,你得跑完 17 小时才隐约觉得这模型怎么老在复读,然后 68 卡时的钱已经花掉了。

五个步骤各自解决一个具体问题:

Step做什么不做会怎样
1 下载把语料弄到本地
2 清洗扔掉太短的、结构性垃圾、语种不对的模型学一堆导航栏和分隔线
3 去重扔掉重复和近似重复的文档重复内容被反复加权,模型爱复读
4 训 tokenizer造出那张「片段 ↔ 编号」的对照表用现成的表,中文压缩率极差
5 打包tokenize 后拼成一个大数组存盘训练时 CPU 解析拖死 GPU

这五步顺序不能换,原因在 3.6 节和 4.6 节会讲。

整条流水线:

最后产出两个文件:

文件内容估算体积
train.bin所有 token id 首尾接起来的扁平数组约 20 GB
val.bin同上,留出来的验证集约 20 MB

20 GB 是这么来的:10B 个 token,每个 token 存 2 字节,就是 20 GB。为什么一个 token 只要 2 字节,第四节讲词表大小的时候说。

一、Step 1:把语料下下来

1.1 先下小的那个

不要一上来就拉 28.5 GB。先用一个很小的数据集把整条流水线跑通,确认代码没问题,再换大的。

我用 roneneldan/TinyStories,1.0 GB、2,141,709 行,是一批用简单词汇写的儿童故事。几分钟能跑完一轮。它训出来的模型很蠢,但足够验证代码是对的。

# code/download_data.py
from datasets import load_dataset

# 冒烟测试用:TinyStories,1.0 GB
ds = load_dataset("roneneldan/TinyStories", split="train")
ds.to_parquet("data/raw/tinystories.parquet")

# 正式用:FineWeb-Edu 的 10B token 子集,28.5 GB
ds = load_dataset(
"HuggingFaceFW/fineweb-edu",
name="sample-10BT",
split="train",
num_proc=16, # 下载和解析都吃 CPU,开满
)
ds.to_parquet("data/raw/fineweb_edu_10bt.parquet")

这里的 .parquet 是一种列式存储格式,可以理解成「压缩过的、按列组织的 csv」。比 csv 和 json 小很多,读起来也快,是数据集分发的事实标准。

国内直连 HuggingFace 通常不通,用镜像站:

export HF_ENDPOINT=https://hf-mirror.com

1.2 下载必须在租卡之前做完

00 篇提过一次,这里再说一遍,因为这是最容易亏钱的地方。

下载 28.5 GB 按 100 Mbps 算要 40 分钟以上,实际经常更久。要是开着 4 张 A100 在那等下载,这 40 分钟是按 GPU 单价付的钱,纯浪费。

正确顺序:先在便宜的 CPU 机器或者本地把 Step 1 到 Step 5 全做完,把最终的 train.bin 传到持久化云盘,然后才开 GPU 实例,挂载云盘直接开训。

整个 01 篇都不需要 GPU。

1.3 下完先看几眼

别急着往下走,先打开看看。这一步经常能提前发现问题。

from datasets import load_dataset

ds = load_dataset("parquet", data_files="data/raw/fineweb_edu_10bt.parquet", split="train")
print(ds) # 看行数和字段名
print(ds[0].keys()) # 看有哪些列
for i in range(3):
print("=" * 60)
print(ds[i]["text"][:500])

要确认的事:正文在哪个字段(FineWeb-Edu 是 text)、有没有 HTML 标签残留、有没有明显乱码、文档大概多长。

二、Step 2:清洗

2.1 这一步解决什么问题

网页正文抽取是不可能百分百干净的。抽出来的文本里经常混着导航栏、页脚、版权声明、分隔线,还有一些只有标题没有正文的空壳页面。

这些东西喂给模型,模型就会认真地学它们。学出来的结果就是模型时不时冒出一句「版权所有 转载请注明出处」。

FineWeb-Edu 已经是筛过一轮的高质量子集,是用一个教育质量分类器从 FineWeb 里挑出来的,重活别人干完了。所以我们这层清洗是补漏,规则要保守。

2.2 规则和阈值

规则阈值为什么是这个阈值
文档太短< 200 字符短于这个基本是标题、菜单项,没有上下文可学
重复行占比> 30%同一行重复多次,典型是导航栏或页脚残留
单字符连续重复> 50====== 这类分隔线,以及爬虫产生的垃圾
非字母数字占比> 50%代码块、表格、乱码
语种置信度< 0.8判不准的多半是混合语种或乱码

这些阈值没有理论依据,都是经验值,来自 C4、Gopher、FineWeb 这几套公开的清洗方案。真正重要的不是阈值本身,是下一节那件事。

2.3 每条规则都必须打点计数

这是这一节最要紧的一句话。每条规则删了多少条,都要记下来,跑完打一张表。

# code/clean.py
import re
from collections import Counter

stats = Counter()

def is_clean(text: str) -> bool:
if len(text) < 200:
stats["too_short"] += 1
return False

lines = text.split("\n")
if lines:
dup_ratio = 1 - len(set(lines)) / len(lines)
if dup_ratio > 0.30:
stats["dup_lines"] += 1
return False

if re.search(r"(.)\1{50,}", text):
stats["char_repeat"] += 1
return False

alnum = sum(c.isalnum() for c in text)
if alnum / len(text) < 0.50:
stats["low_alnum"] += 1
return False

stats["kept"] += 1
return True


def report(total):
print(f"{'规则':<14}{'删除条数':>12}{'占比':>10}")
for k, v in stats.most_common():
if k == "kept":
continue
print(f"{k:<14}{v:>12,}{v / total * 100:>9.2f}%")
print(f"{'保留':<14}{stats['kept']:>12,}{stats['kept'] / total * 100:>9.2f}%")

为什么非要打点。因为清洗规则写错了是静默失败,不报错,只是数据悄悄少了一块。

举个真会踩的例子。上面 low_alnum 那条用了 c.isalnum()。Python 里汉字调用 isalnum() 返回的是 True,所以中文语料不受影响。但如果你手滑写成 c.isascii() and c.isalnum(),中文的 alnum 占比就变成 0,整个中文语料会被一条不起眼的规则全部删光,而程序不会报任何错。

有统计表,你会立刻看到 low_alnum 那行占比 99%,两秒钟发现问题。没有统计表,你要等到模型训完发现不会说中文。

这张表跑完要存下来,05 篇复盘要用。

三、Step 3:去重

3.1 重复数据为什么危险

同一段文本在语料里出现 N 次,模型就会在它上面走 N 次梯度,等于给这段文本加权了 N 倍。

后果有三个。模型在这些片段上过拟合,生成的时候容易原样背出来。有效数据量虚高,你以为训了 10B token,实际独立信息可能只有 8B。验证集如果和训练集有重叠,验证 loss 会偏低,让你误判模型比实际更好。

网页语料的重复率相当高,转载、镜像站、模板化页面都会产生大量重复。去重是性价比最高的一步清洗。

要去两类重复,方法完全不同,两类都得处理:

  • 完全相同:一模一样的转载页。用哈希,简单。
  • 高度相似:换了个标题、改了两个词。用 MinHash,是这一节的主要内容。

3.2 第一类:完全相同的,用哈希

# code/dedup_exact.py
import hashlib

def exact_dedup(docs):
seen = set()
for doc in docs:
# 先归一化,避免只差空白字符的被当成不同文档
norm = " ".join(doc.split()).lower()
h = hashlib.sha256(norm.encode("utf-8")).digest()
if h in seen:
continue
seen.add(h)
yield doc

两个细节。一是先归一化再算哈希,否则「多了一个空格」就被当成不同文档,这一步就白做了。二是存 digest()(32 字节的 bytes)而不是 hexdigest()(64 字符的字符串),内存差一倍。1000 万篇文档大约要 1 GB 内存,还能接受;再大就要上布隆过滤器或者分片处理。

3.3 第二类:高度相似的,问题在哪

哈希抓不到「只改了两个字」的情况,因为哈希值会完全变掉。

那能不能直接两两比较?不行。1000 万篇文档两两比是 5×10¹³ 次比较,每次比较还要遍历两篇文档的全文。这个量级跑不完。

MinHash 要解决的就是这个:怎么把「比较两篇文档」这件事变得极快。

思路分两层,先把每篇文档压成一个很短的签名(3.4 节),再想办法避免两两比较(3.5 节)。

3.4 MinHash:把文档压成定长签名

第一步,把文档变成一个集合。

用滑动窗口把文本切成一小段一小段,每段长度固定(一般 5 个字符),这些片段叫 shingle。切完去重,得到一个集合。

def shingles(text, n=5):
return {text[i:i + n] for i in range(len(text) - n + 1)}

比如 "abcdef" 切成 {"abcde", "bcdef"}

变成集合之后,两篇文档的相似度就可以用 Jaccard 相似度定义,就是交集除以并集:

J(A,B)=ABABJ(A, B) = \frac{|A \cap B|}{|A \cup B|}

完全相同是 1,完全没有共同片段是 0。

第二步,用「取最小值」这个技巧压缩。

现在集合还是很大,直接存不划算。MinHash 的做法是:

准备一个哈希函数,把集合里每个 shingle 都哈希一遍,只保留最小的那个哈希值。整个集合就被压成了一个数字。

用 k 个不同的哈希函数各做一次,就得到 k 个数字。这 k 个数字排起来就是这篇文档的签名,长度固定为 k,跟原文多长无关。

第三步,为什么这样能估相似度。

这是 MinHash 最巧妙的地方。结论是:

对任意一个哈希函数,两篇文档的最小哈希值相同的概率,恰好等于它们的 Jaccard 相似度。

推导只要一句话。把 ABA \cup B 里所有元素按哈希值从小到大排序,问「排第一的那个元素属于 ABA \cap B 的概率是多少」。因为哈希值可以看作随机的,这个概率就是 ABAB\frac{|A \cap B|}{|A \cup B|},正好是 Jaccard。而只有排第一的元素同时属于两个集合,A 和 B 的最小值才会相等。

所以拿 k 个哈希函数的签名一比,相同位的比例就是 Jaccard 的估计值。k 越大越准。

第四步,实测一下到底准不准。

原理讲完了要验。拿两句只差两个字的中文,和一句完全无关的:

A = "机器学习模型的训练过程需要大量的高质量语料数据来支撑最终的效果表现"
B = "机器学习模型的训练过程需要大量的高品质语料数据来支撑最终的效果表现"
C = "今天天气不错适合出门散步顺便去公园里看看新开的那些花草树木"

真实 Jaccard:AB = 0.6571,AC = 0.0000。MinHash 的估计值:

k估计 A~B估计 A~C
160.56250.0000
640.64060.0000
1280.65620.0000
2560.68360.0000

(这组数是跑出来的,脚本在 code/minhash_demo.py。)

k=128 时估出来 0.6562,真值 0.6571,差了不到千分之一。而签名只有 128 个整数,不管原文是 200 字还是 20 万字。这就是 MinHash 的价值。

3.5 LSH 分桶:避免两两比较

签名短了,但如果还是两两比,1000 万篇还是 5×10¹³ 次,还是跑不完。

LSH(局部敏感哈希)的思路:让相似的文档「自动碰撞到同一个桶里」,只在桶内比较。

做法是把长度 k 的签名切成 b 段,每段 r 行(k=b×rk = b \times r)。每一段单独算个哈希,扔进对应的桶。只要两篇文档有任意一段完全相同,就落进同一个桶,成为候选对。然后只对候选对做精确比较。

相似度为 s 的两篇文档成为候选对的概率是:

P=1(1sr)bP = 1 - (1 - s^r)^b

推导很直白。某一段的 r 行全都相同的概率是 srs^r。这一段不同的概率是 1sr1 - s^r。b 段全都不同的概率是 (1sr)b(1 - s^r)^b。至少有一段相同,就是 1 减去它。

这个公式画出来是条 S 形曲线。调 b 和 r 就是在调曲线的陡峭位置,也就是在调「相似到什么程度才算重复」。k=128 时几组配置的实测值:

相似度 sb=8, r=16b=16, r=8b=32, r=4
0.300.0000.0010.229
0.500.0000.0610.873
0.600.0020.2370.988
0.700.0260.6131.000
0.800.2040.9471.000
0.900.8061.0001.000
0.950.9901.0001.000

怎么读这张表。

b=32, r=4 太松:相似度只有 0.3 的文档都有 22.9% 概率进候选桶,会产生海量候选对,后面的精确比较会被拖死。

b=8, r=16 太严:相似度 0.8 的文档只有 20.4% 概率被抓到,大量近似重复漏过去了。

b=16, r=8 平衡:0.8 以上抓得很牢(94.7%),0.5 以下基本放过(6.1%)。去重阈值定在 0.8 附近的话,这组参数就合适。

3.6 实际动手:用 datatrove

上面的原理自己写一遍很有价值,弄懂了才知道参数在调什么。但真要处理 28.5 GB,用 HuggingFace 的 datatrove。FineWeb 本身就是拿它处理的,参数经过大规模验证,而且分片、落盘、断点续跑都做好了。

# code/dedup_minhash.py
from datatrove.pipeline.dedup import MinhashDedupSignature
from datatrove.pipeline.dedup.minhash import MinhashConfig
from datatrove.executor import LocalPipelineExecutor
from datatrove.pipeline.readers import ParquetReader

config = MinhashConfig(
num_buckets=16, # 就是上面的 b
hashes_per_bucket=8, # 就是上面的 r,b*r=128
n_grams=5,
)

executor = LocalPipelineExecutor(
pipeline=[
ParquetReader("data/cleaned/"),
MinhashDedupSignature(output_folder="data/minhash/sigs", config=config),
],
tasks=32, # 按 CPU 核数调
)
executor.run()

datatrove 的 MinHash 去重要分四个阶段依次跑:算签名、按桶找候选对、把候选对聚成连通分量、按聚类结果过滤。中间结果落盘,所以中断了可以接着跑。完整脚本在 code/dedup_minhash.py

3.7 为什么清洗要在去重之前

清洗会改变文本。两篇文档原本因为页脚不同而被判为不同,把页脚去掉之后就变成一样的了,这时候去重才能抓到。

反过来先去重再清洗,这批重复就漏掉了。

四、Step 4:训 Tokenizer

回到 0.2 节留下的问题:怎么决定哪些片段进词表。

4.1 BPE 算法:核心就一句话

我们用的算法叫 BPE(Byte Pair Encoding),原本是个数据压缩算法。规则一句话说完:

反复找出现频次最高的相邻符号对,把它合并成一个新符号,重复 N 次。

初始状态是把所有词拆成单个字符。每合并一次,词表就多一个条目。合并 32000 次左右,词表就建好了。

4.2 手推一遍

拿一个只有四个词的小语料推一遍,看它到底怎么工作。词和频次:

low: 5    lower: 2    newest: 6    widest: 3

先把每个词拆成字符,词尾加个 </w> 标记词的边界(不然 lowlower 里的 low 分不清是不是完整的词)。然后开始迭代。

下面是实际跑出来的输出,脚本在 code/bpe_demo.py

--- 第 1 次合并 ---
频次最高的几对: (e+s):9, (s+t):9, (t+</w>):9, (w+e):8
合并: e + s -> es (出现 9 次)
结果: l o w </w> ×5 | l o w e r </w> ×2 | n e w es t </w> ×6 | w i d es t </w> ×3

--- 第 2 次合并 ---
频次最高的几对: (es+t):9, (t+</w>):9, (l+o):7, (o+w):7
合并: es + t -> est (出现 9 次)
结果: l o w </w> ×5 | l o w e r </w> ×2 | n e w est </w> ×6 | w i d est </w> ×3

--- 第 3 次合并 ---
频次最高的几对: (est+</w>):9, (l+o):7, (o+w):7, (n+e):6
合并: est + </w> -> est</w> (出现 9 次)
结果: l o w </w> ×5 | l o w e r </w> ×2 | n e w est</w> ×6 | w i d est</w> ×3

--- 第 4 次合并 ---
频次最高的几对: (l+o):7, (o+w):7, (n+e):6, (e+w):6
合并: l + o -> lo (出现 7 次)
结果: lo w </w> ×5 | lo w e r </w> ×2 | n e w est</w> ×6 | w i d est</w> ×3

--- 第 5 次合并 ---
频次最高的几对: (lo+w):7, (n+e):6, (e+w):6, (w+est</w>):6
合并: lo + w -> low (出现 7 次)
结果: low </w> ×5 | low e r </w> ×2 | ne w est</w> ×6 | w i d est</w> ×3

第 1 到 3 次合并,算法自己发现了 est</w> 这个后缀。它完全不知道什么叫「形容词最高级」,只是因为 newestwidest 都有这个模式,频次加起来 9,最高。第 4、5 次合出了 low,因为 lowlower 共享词根。

这就是 BPE 的全部机制。没有任何语言学知识,纯粹统计。跑几万次之后,高频词会变成完整的一个 token,低频词自然被拆成片段。

4.3 顺带说一下另外两种

子词算法主流有三种,区别只在「按什么标准选择合并哪一对」。

BPE 按频次,选出现次数最多的那对。Llama、GPT 系列用它。

WordPiece 按似然增益,选合并后能让语言模型似然提升最大的那对,判据大致是 P(ab)P(a)P(b)\frac{P(ab)}{P(a)P(b)}。它更偏好「两个片段总是一起出现」的组合,而不单看绝对频次。BERT 用它。

Unigram 反着来,先给一个很大的初始词表,迭代删掉「删了对总似然影响最小」的子词,直到降到目标大小。T5、mT5 用它。

我们选 BPE,理由是对齐 Llama,出问题时方便和现成实现对照。

4.4 Byte-level:让词表永远不会「不认识」

直接对字符做 BPE 有个漏洞。如果推理时遇到训练语料里从没出现过的字符(某个生僻字、某个新 emoji),词表里查不到,只能标记成「未知」,这个位置的信息就丢了。

解决办法是不在字符层面做 BPE,而在 UTF-8 字节层面做。

任何文本编码成 UTF-8 之后,都是一串 0 到 255 的字节。所以只要初始词表包含全部 256 个字节,任何输入都一定能被编码,永远不会出现「不认识」的情况

这就是后面代码里 initial_alphabet=pre_tokenizers.ByteLevel.alphabet() 那一行的作用。这行漏了,就丢掉了 byte-level BPE 最大的好处。

4.5 压缩率:中文为什么吃亏

压缩率指的是平均一个 token 能装下多少原始字符。它直接决定了同样 10B 个 token 里到底装了多少真信息。

中英文差别很大,根子在 UTF-8 编码:

字符UTF-8 字节数
英文字母 a1
汉字 3
中文标点 3
emoji 😀4

对比一下:

"machine learning model training"   31 字符 / 31 字节  → 1.00 字节每字符
"机器学习模型训练过程" 10 字符 / 30 字节 → 3.00 字节每字符

byte-level BPE 在字节上跑,所以一个汉字起步就是 3 个字节。如果 tokenizer 的训练语料里中文占比很低,这 3 个字节压根没机会被合并起来,一个汉字就要吃掉 2 到 3 个 token。同样 10B 的预算,中文语料装进去的实际信息量可能只有英文的三分之一。

结论:如果模型要处理中文,tokenizer 的训练语料里就必须有足够比例的中文,让 BPE 有机会把常见汉字和词组合并成单个 token。这是自训 tokenizer 相比直接拿现成英文 tokenizer 的主要收益。

三组对照的实测数字我先留空,跑完回填。计划测 GPT-2 的 50257 词表、Llama 的 32000 词表、自训的 32000 词表,各自在英文和中文语料上的「字符 / token」比值。不实测就写数字,是这个专题明确不做的事。

4.6 词表大小定 32000

回答 0.5 节留下的「为什么一个 token 只要 2 字节」。

一个 token id 要存下来,占几个字节取决于最大的 id 是多少。uint16 能表示 0 到 65535。我们的词表是 32000,所有 id 都小于 65535,所以每个 token 存 2 字节就够。

要是贪心把词表开到 65536 以上,就得用 uint32,每个 token 4 字节,train.bin 从 20 GB 变成 40 GB,训练时读盘带宽压力直接翻倍。

另一个理由是对齐 Llama,出问题时方便对照。

顺带说词表大小的权衡。词表越大,压缩率越好,序列越短,同样上下文能装更多信息。但 embedding 层和输出层的参数量正比于词表大小,词表翻倍这两层就翻倍。我们这个配置里 embedding 已经占了总参数的 9.8%(49M / 502M),要是词表开到 128k,光 embedding 就是 196M,占总量三成多,那就本末倒置了。小模型配小词表,是有道理的。

4.7 动手训,分五步

Step 1,装依赖。

pip install tokenizers datasets

Step 2,准备语料迭代器。

不需要全量,采样 2 到 5 GB 就够训出稳定的合并规则。注意采样要保持中英文比例和最终训练语料一致,否则压缩率会跑偏。

# code/train_tokenizer.py
from datasets import load_dataset

def corpus_iter(paths, sample_ratio=0.05, batch_size=1000):
ds = load_dataset("parquet", data_files=paths, split="train", streaming=True)
buf = []
for i, row in enumerate(ds):
if i % int(1 / sample_ratio) != 0:
continue
buf.append(row["text"])
if len(buf) >= batch_size:
yield buf
buf = []
if buf:
yield buf

streaming=True 是为了不把整个数据集读进内存。按 batch 产出是因为 tokenizers 库批量处理比逐条快很多。

Step 3,配置 tokenizer。

from tokenizers import Tokenizer, models, trainers, pre_tokenizers, decoders

tokenizer = Tokenizer(models.BPE())
tokenizer.pre_tokenizer = pre_tokenizers.ByteLevel(add_prefix_space=False)
tokenizer.decoder = decoders.ByteLevel()

pre_tokenizer 是在跑 BPE 之前先做一次粗切分,ByteLevel 会把文本转成字节并做一层可逆映射。decoder 必须和它配套,否则 decode 回来是乱码。这两行要么都用 ByteLevel,要么都不用,不能混搭。

Step 4,配训练器,开训。

trainer = trainers.BpeTrainer(
vocab_size=32000,
special_tokens=["<|endoftext|>", "<|im_start|>", "<|im_end|>"],
initial_alphabet=pre_tokenizers.ByteLevel.alphabet(), # 256 个字节全给上
show_progress=True,
)

tokenizer.train_from_iterator(corpus_iter(["data/cleaned/*.parquet"]), trainer=trainer)
tokenizer.save("tokenizer.json")

Step 5,验证。 训完一定要验,别直接进下一步。

tok = Tokenizer.from_file("tokenizer.json")

# 1. 词表大小对不对
assert tok.get_vocab_size() == 32000

# 2. 编码再解码,能不能一字不差地还原
for s in ["Hello world", "机器学习模型训练", "emoji 😀 测试", "a" * 200]:
ids = tok.encode(s).ids
back = tok.decode(ids)
assert back == s, f"还原失败: {s!r} -> {back!r}"
print(f"{len(s):>4} 字符 -> {len(ids):>4} token 压缩率 {len(s) / len(ids):.2f}")

# 3. special token 的 id 拿得到
for t in ["<|endoftext|>", "<|im_start|>", "<|im_end|>"]:
print(t, tok.token_to_id(t))

第 2 条最重要,叫无损往返。编码再解码必须一字不差还原原文。还原不了说明 pre_tokenizerdecoder 没配套,或者 initial_alphabet 漏了字节。这时候如果继续往下走,训出来的模型会一直带着这个 bug。emoji 那条是专门测多字节字符的。

4.8 特殊 token 现在就得留好

<|im_start|><|im_end|> 是用来标记对话角色的,预训练根本用不上,但 06 篇做 SFT 要靠它们。

为什么不能到时候再加?因为词表变了,embedding 层的形状就变了(32000 行变成 32002 行),预训练学到的 embedding 权重对不上,只能重训。

现在多花一分钟在 special_tokens 里多写两个,能省掉后面 17 小时。

建议再多留几个占位符,比如 <|reserved_0|><|reserved_7|>,将来想加工具调用标记、思维链标记之类的都不用动词表。Llama 3 就是这么干的。

五、Step 5:打包成 bin

5.1 为什么最后要存成一个大数组

现在数据已经干净了,tokenizer 也有了。最后一步是把所有文档 tokenize,拼成一个大数组存盘。

为什么不能训练的时候现场 tokenize?算笔账就明白了。

按 04 篇的配置,4 张 A100、batch size 512、序列长 2048,一步就要 1,048,576 个 token。整个训练要走几十万步。如果每一步都现场解析 parquet 再 tokenize,光这部分就要几百毫秒,而一步前向反向也就几百毫秒。结果是 GPU 有一半时间在等 CPU,00 篇算的 MFU 直接腰斩,等于白租两张卡。

所以要提前把最重的活干完,训练时只剩「从数组里取一段」这一个动作。

5.2 memmap:文件比内存大也能读

train.bin 有 20 GB,多数机器的内存装不下。这里用 np.memmap 解决。

memmap 的原理是把文件映射到虚拟地址空间,不真的读进内存,实际读哪一页由操作系统的 page cache 管。随机读几个 batch,操作系统只把碰到的那几页调进来。所以哪怕文件比内存大得多,也能正常跑。

用起来和普通 numpy 数组一模一样:

import numpy as np

data = np.memmap("train.bin", dtype=np.uint16, mode="r")

def get_batch(batch_size, seq_len):
# 在整个 token 流里随机取起点
ix = np.random.randint(len(data) - seq_len - 1, size=batch_size)
x = np.stack([data[i : i + seq_len].astype(np.int64) for i in ix])
y = np.stack([data[i + 1 : i + 1 + seq_len].astype(np.int64) for i in ix])
return x, y

对照 0.3 节:y 就是 x 右移一位,和那里画的一模一样,只是一次取 batch_size 条。

5.3 为什么可以不管文档边界

上面取起点是在整个 token 流里随便取的,可能落在某篇文档中间,甚至横跨两篇。这在预训练阶段是标准做法,两个理由。

一是不用 padding。 如果按文档切,文档长短不一,短的要补占位 token 对齐长度,而占位的位置是不算 loss 的,那部分算力就白花了。语料里短文档占比不低,浪费可能到两三成。首尾拼接之后每个 token 都在贡献梯度,一点不浪费。

二是模型自己能学会边界。 文档之间插了 EOS token(end of sequence,标记「这篇结束了」的特殊 token),模型见多了就知道 EOS 之后是新话题,不该再往前依赖。这个能力它自己从数据里学,不用我们显式告诉它。

代价是极少数样本跨了文档边界,前半截和后半截没关系。在 10B token 这个量级下,这点噪声可以忽略。

到 SFT 阶段就不能这么干了,那时候要精确控制只在回复上算 loss,得按样本组织、得 padding。这个区别 06 篇讲。

5.4 基本版打包代码

# code/pack.py
import numpy as np
from tokenizers import Tokenizer

tok = Tokenizer.from_file("tokenizer.json")
EOS = tok.token_to_id("<|endoftext|>")

def pack(docs, out_path, total_tokens_estimate):
# 预分配,避免边写边扩容
arr = np.memmap(out_path, dtype=np.uint16, mode="w+",
shape=(total_tokens_estimate,))
idx = 0
for doc in docs:
ids = tok.encode(doc).ids
ids.append(EOS) # 每篇末尾插一个 EOS
arr[idx : idx + len(ids)] = ids
idx += len(ids)
arr.flush()
return idx # 真实写入量,拿它去截断文件

5.5 必须并行

单进程 tokenize 10B token 是按天算的,必须并行。做法是按 shard(分片)切分,每个进程处理一个 shard 写自己的文件,最后拼接。

# code/pack_parallel.py
import numpy as np, os
from multiprocessing import Pool
from tokenizers import Tokenizer

tok = None

def init_worker():
global tok
tok = Tokenizer.from_file("tokenizer.json") # 每个进程各自加载

def encode_shard(args):
shard_id, texts, out_dir = args
eos = tok.token_to_id("<|endoftext|>")
ids = []
for t in texts:
ids.extend(tok.encode(t).ids)
ids.append(eos)
arr = np.array(ids, dtype=np.uint16)
path = os.path.join(out_dir, f"shard_{shard_id:05d}.bin")
arr.tofile(path)
return path, len(arr)

def run(shards, out_dir, workers=32):
os.makedirs(out_dir, exist_ok=True)
with Pool(workers, initializer=init_worker) as pool:
results = pool.map(encode_shard,
[(i, s, out_dir) for i, s in enumerate(shards)])
total = sum(n for _, n in results)
print(f"{len(results)} 个 shard,共 {total:,} tokens")
return results

这里有个坑。Tokenizer 对象不能直接跨进程传递,必须在每个 worker 里各自 from_file 加载,也就是 init_worker 干的事。如果图省事把 tok 当全局变量让子进程继承,在 Linux 上碰巧能跑(fork 会复制内存),但在 macOS 和 Windows 上会直接报 pickle 错误。写成 initializer 三个平台都稳。

拼接直接按二进制拼,因为都是同一种 dtype 的扁平数组:

def concat(shard_paths, out_path):
with open(out_path, "wb") as fout:
for p in sorted(shard_paths):
with open(p, "rb") as fin:
while chunk := fin.read(1 << 24): # 16 MB 一块
fout.write(chunk)

5.6 容量怎么估

np.memmapmode="w+" 时必须预先给定长度,写超了直接崩。所以要先估一下总共有多少 token。

做法是在 1% 的样本上跑一遍,量出平均「字节 / token」,再乘回全量,最后上浮 5% 留余量:

def estimate_tokens(sample_texts, total_bytes):
sample_bytes = sum(len(t.encode("utf-8")) for t in sample_texts)
sample_tokens = sum(len(tok.encode(t).ids) + 1 for t in sample_texts)
bytes_per_token = sample_bytes / sample_tokens
return int(total_bytes / bytes_per_token * 1.05)

用 5.5 节的并行 shard 方案其实不需要估,因为每个 shard 是 np.array(...).tofile() 直接写的,长度自己知道。只有单进程 memmap 方案才要预估。

5.7 验证集从源头切开

在去重之后、打包之前,随机抽 0.1% 的文档当验证集。

关键是抽文档,不是抽 token。如果在 token 层面切,同一篇文档会同时出现在 train 和 val 里,验证 loss 就废了,会明显偏低,让你误以为模型效果很好。

import random

def split_train_val(docs, val_ratio=0.001, seed=42):
docs = list(docs)
random.Random(seed).shuffle(docs)
n_val = max(1, int(len(docs) * val_ratio))
return docs[n_val:], docs[:n_val]

固定 seed,这样重跑能复现同样的切分。

六、验收清单

这一篇做完,下面每条都要能答上来,不然别进 02 篇。

  • train.binval.bin 已生成,体积和预估对得上
  • 随机取 10 个位置 decode 回文本,读着是通顺的
  • tokenizer 编解码往返无损,中英文和 emoji 都测过
  • EOS 在 token 流里的密度合理,大致等于 1 除以平均文档 token 数
  • 清洗各条规则的删除量已打表,没有哪条高得离谱
  • 去重前后的文档数和体积都记下来了
  • tokenizer 压缩率已实测,中英文分别有数
  • 确认 val.bin 里的文档没出现在 train.bin
  • 用 TinyStories 走完了同一条流水线

前几条的验证代码:

data = np.memmap("train.bin", dtype=np.uint16, mode="r")
print(f"总 token 数: {len(data):,}")

eos_count = int((np.asarray(data[:10_000_000]) == EOS).sum())
print(f"前 1000 万 token 里有 {eos_count:,} 个 EOS,"
f"平均文档长度约 {10_000_000 // max(eos_count, 1):,} token")

for start in np.random.randint(0, len(data) - 200, size=3):
print("=" * 60)
print(tok.decode([int(x) for x in data[start:start + 200]]))

最后那段 decode 出来的文本如果读着通顺,说明整条流水线是对的。如果是乱码,回 4.7 节 Step 5 查往返测试。

七、常见问题

现象大概率是什么原因
decode 回来是乱码pre_tokenizerdecoder 没配套,检查是不是都用了 ByteLevel
某个字符编码后变成未知 tokeninitial_alphabet 没给全 256 个字节
压缩率异常低,中文一个字占好几个 tokentokenizer 训练语料里中文占比太少
某条清洗规则删掉了大半数据规则写错了,尤其检查中文场景下的 isalnum 之类判断
memmap 写到一半崩了预估的 total_tokens_estimate 不够,上浮比例调大
多进程 tokenize 报 pickle 错误Tokenizer 对象跨进程传了,改成在 initializer 里各自加载
验证 loss 明显低于训练 lossval 是在 token 层面切的,和 train 有重叠,改成按文档切
先在 TinyStories 上跑通

roneneldan/TinyStories 只有 1.0 GB、2,141,709 行,整条流水线几分钟跑完一轮。上面九条先在它身上全部验收通过,再去碰 28.5 GB 的 FineWeb-Edu。小数据上发现的 bug,代价是几分钟;大数据上发现的,代价是几小时加钱。